from sqlalchemy import Column, Integer, String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from init.mysql import engine

Base = declarative_base()


class StockProfit(Base):
    __tablename__ = 'stock_profit'
    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(String(10))
    date = Column(String(10))
    income = Column(String(20))
    expend = Column(String(30))
    profit = Column(String(20))
    totalp = Column(String(20))
    reprofit = Column(String(20))
    basege = Column(String(20))
    ettege = Column(String(20))
    otherp = Column(String(20))
    totalcp = Column(String(20))

    def __init__(self, code, date, income, expend, profit, totalp, reprofit, basege, ettege, otherp, totalcp):
        self.code = code
        self.date = date
        self.income = income
        self.expend = expend
        self.profit = profit
        self.totalp = totalp
        self.reprofit = reprofit
        self.basege = basege
        self.ettege = ettege
        self.otherp = otherp
        self.totalcp = totalcp

    @staticmethod
    def insert_or_update(code, data_list):
        session = sessionmaker(bind=engine)()
        for data in data_list:
            stock_profit = session.query(StockProfit).filter_by(code=code, date=data['date']).first()
            if stock_profit:
                stock_profit.code = code
                stock_profit.date = data['date']
                stock_profit.income = data['income']
                stock_profit.expend = data['expend']
                stock_profit.profit = data['profit']
                stock_profit.totalp = data['totalp']
                stock_profit.reprofit = data['reprofit']
                stock_profit.basege = data['basege']
                stock_profit.ettege = data['ettege']
                stock_profit.otherp = data['otherp']
                stock_profit.totalcp = data['totalcp']
            else:
                session.add(StockProfit(code=code,date=data['date'], income=data['income'], expend=data['expend'],
                                        profit=data['profit'], totalp=data['totalp'], reprofit=data['reprofit'],
                                        basege=data['basege'], ettege=data['ettege'], otherp=data['otherp'],
                                        totalcp=data['totalcp']))
        session.commit()


Base.metadata.create_all(engine)
